without @tf.function OOM

def train_step(self, x, step, grad_clip=True, clip_value=1.0):
    print('tracing', x, step, grad_clip, clip_value)
    with tf.name_scope("input_data"):
      ...